#include "kernel_operator.h"
#include <typeinfo>
#include <type_traits>
using namespace AscendC;

constexpr int32_t BUFFER_NUM = 2;      //昇腾双buffer技术
     
template<typename TYPE_X1,typename TYPE_X2,typename TYPE_Y> class KernelIsClose {    
using T = TYPE_X1;
public:
    __aicore__ inline KernelIsClose() {}

    __aicore__ inline void serial_version()
    {
        /*串行版本*/
        for(int i=0; i< this->coreDataNum;++i)
        {
            float input = 0; float other = 0;
            if constexpr (std::is_same_v<T, uint8_t>) {
                input = static_cast<float>((int)xGm1.GetValue(i));
                other = static_cast<float>((int)xGm2.GetValue(i));
            }
            else{ 
                printf("...not uint8_t...");
                input = (float)xGm1.GetValue(i);
                other = (float)xGm2.GetValue(i);
                printf("input is  %d\n",isnan(input));
                printf("other is %d\n",isnan(other));
                if((this->equal_nan == true) && (isnan(input) && isnan(other)))
                {
                    yGm.SetValue(i,1);
                    continue;
                }
                if((isnan(input)) || isnan(other))
                {
                    yGm.SetValue(i,0);
                    continue;
                }
                
            }

            //1、input = |input - other| 
            if(input > other)
                input = input - other;
            else
                input = other - input;
            //printf("input is %f",input);
            //3、other = atol + rtol * |other|
            if (other > 0)
                other = this->atol + this->rtol * other ;
            else
                other =  this->atol - this->rtol * other;
            //printf(",other is %f\n",other);
            //4、compare 
            if(input <= other)
                yGm.SetValue(i,1);
            else
                yGm.SetValue(i,0);
            //printf("yGm.GetValue is %f\n",yGm.GetValue(i));      
        }
    }


    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y,  uint32_t CoreDataNum, uint32_t finalTileNum, uint32_t tileDataNum, uint32_t TailDataNum, float rtol, float atol, bool equal_nan)
    {

        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        //初始化变量
        this->coreDataNum = CoreDataNum;
        this->tileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->tailDataNum = TailDataNum;
        this->rtol = rtol;
        this->atol = atol;
        this->equal_nan = equal_nan;
        //printf("CoreDataNum %d,finalTileNum %d,tileDataNum %d, TailDataNum %d\n",CoreDataNum,finalTileNum,tileDataNum,TailDataNum );

        //初始化globalbuffer和localbuffer
        xGm1.SetGlobalBuffer((__gm__ DTYPE_X1*)x1, this->coreDataNum);
        xGm2.SetGlobalBuffer((__gm__ DTYPE_X2*)x2, this->coreDataNum);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, this->coreDataNum);

        pipe.InitBuffer(inQueueX1, BUFFER_NUM, (this->tileDataNum+ 256)  * sizeof(DTYPE_X1));  
        pipe.InitBuffer(inQueueX2, BUFFER_NUM, (this->tileDataNum+ 256)  * sizeof(DTYPE_X2));  
        pipe.InitBuffer(outQueueY, BUFFER_NUM, (this->tileDataNum+ 256)  * sizeof(DTYPE_Y));     
        if constexpr (std::is_same_v<DTYPE_X1, int32_t>)
        {
            //printf("1  ");
            pipe.InitBuffer(QueueFloat1, (this->tileDataNum + 256)  * sizeof(float));
            pipe.InitBuffer(QueueFloat2, (this->tileDataNum + 256)  * sizeof(float));   
        }
        else if constexpr (std::is_same_v<DTYPE_X1, uint8_t>)
        {
            //printf("2  ");
            pipe.InitBuffer(QueueHalf1, (this->tileDataNum + 256)  * sizeof(half));
            pipe.InitBuffer(QueueHalf2, (this->tileDataNum + 256)  * sizeof(half));
        }
        else if constexpr (std::is_same_v<DTYPE_X1, half>)
        {
            //printf("3  ");
            pipe.InitBuffer(QueueFloat1, (this->tileDataNum + 256)  * sizeof(float));
            pipe.InitBuffer(QueueFloat2, (this->tileDataNum + 256)  * sizeof(float)); 
            //printf("this->tileDataNum + 256 is %d",this->tileDataNum + 256);
        }
        else
        {
            //printf("4  ");
            pipe.InitBuffer(QueueFloat1, (this->tileDataNum + 256)  * sizeof(float));
            //pipe.InitBuffer(QueueFloat2, (this->tileDataNum + 256)  * sizeof(float)); 
        }
        pipe.InitBuffer(QueueMask, (this->tileDataNum + 256 )  * sizeof(uint8_t));   //初始化Mask
    }
    __aicore__ inline void Process()
    {
        /*并行版本*/
        // int32_t loopCount = this->tileNum;   //手动考虑了双buff
        // this->processDataNum = this->tileDataNum;
        // for (int32_t i = 0; i < loopCount; i++) {
        //     if (i == this->tileNum - 1) {
        //         this->processDataNum = this->tailDataNum;
        //         //printf(" this->processDataNum is %d",this->processDataNum);
        //     }
        //     CopyIn(i);
        //     Compute(i);
        //     CopyOut(i);
        //     //printf("i finished\n");
        // } 

        /*串行版本*/
        serial_version();

    }

private:

    __aicore__ inline int isnan(float num) 
    {
        unsigned int bits = *(unsigned int*)&num;
        return (bits & 0x7F800000) == 0x7F800000 && (bits & 0x007FFFFF) != 0;
    }

    __aicore__ inline void CopyIn(int32_t progress)
    {
        //考生补充算子代码
        LocalTensor<DTYPE_X1> xLocal1 = inQueueX1.AllocTensor<DTYPE_X1>();
        LocalTensor<DTYPE_X2> xLocal2 = inQueueX2.AllocTensor<DTYPE_X2>();
        DataCopy(xLocal1, xGm1[progress * this->tileDataNum], (this->processDataNum+32)/32*32);  //tileDataNum>tailDataNum，所以不需要麻烦分情况讨论
        DataCopy(xLocal2, xGm2[progress * this->tileDataNum], (this->processDataNum+32)/32*32);
        //printf("xGm1 %f xGm2 %f xLocal1 %f xLocal2 %f\n",xGm1(1), xGm2(1), xLocal1(1), xLocal2(1)); 
        inQueueX1.EnQue(xLocal1);
        inQueueX2.EnQue(xLocal2);
    }
    __aicore__ inline void Compute(int32_t progress)
    {

        LocalTensor<DTYPE_X1> xLocal1 = inQueueX1.DeQue<DTYPE_X1>();
        LocalTensor<DTYPE_X2> xLocal2 = inQueueX2.DeQue<DTYPE_X2>();
        LocalTensor<DTYPE_Y> yLocal   = outQueueY.AllocTensor<DTYPE_Y>();

        auto m_mask    = QueueMask.Get<uint8_t>();

        /****isclose */
        if constexpr (std::is_same_v<DTYPE_X1, uint8_t>)
        {
            //printf( "uint8_t \n ");

            auto m_half1    = QueueHalf1.Get<half>();
            auto m_half2    = QueueHalf2.Get<half>();

            half rtol_half = (half)this->rtol;
            half atol_half = (half)this->atol;
            Cast(m_half1, xLocal1, RoundMode::CAST_NONE, this->processDataNum);
            Cast(m_half2, xLocal2, RoundMode::CAST_NONE, this->processDataNum);

            //1、input = |input - other| 
            Sub(m_half1, m_half1, m_half2, this->processDataNum);
            Abs(m_half1, m_half1, this->processDataNum);
            //2、other =  rtol * |other|
            Abs(m_half2, m_half2, this->processDataNum);
            Muls(m_half2, m_half2, rtol_half, this->processDataNum);
            //3、other = atol + rtol * |other|
            Adds(m_half2, m_half2, atol_half, this->processDataNum);
            //4、compare     
            Compare(m_mask, m_half1, m_half2, CMPMODE::LE, (this->processDataNum+255)/256*256);  //256字节对齐生成mask,小于0为1
            //printf("%d, %d, %d, %d,%d\n",xLocal1(1),xLocal2(1), m_half1(1),m_half2(1),m_mask(1)); 
            //printf("%d, %d, %d, %d,%d\n",xLocal1(2),xLocal2(2), m_half1(2),m_half2(2),m_mask(2)); 
        }
        else if constexpr (std::is_same_v<DTYPE_X1, int32_t>) 
        {
            //printf( "int32_t \n ");

            auto m_float1    = QueueFloat1.Get<float>();
            auto m_float2    = QueueFloat2.Get<float>();
            Cast(m_float1, xLocal1, RoundMode::CAST_NONE, this->processDataNum);
            Cast(m_float2, xLocal2, RoundMode::CAST_NONE, this->processDataNum);

            //1、input = |input - other| 
            Sub(m_float1, m_float1, m_float2, this->processDataNum);
            Abs(m_float1, m_float1, this->processDataNum);
            //2、other =  rtol * |other|
            Abs(m_float2, m_float2, this->processDataNum);
            Muls(m_float2, m_float2, this->rtol, this->processDataNum);
            //3、other = atol + rtol * |other|
            Adds(m_float2, m_float2, this->atol, this->processDataNum);
            //4、compare                
            Compare(m_mask, m_float1, m_float2, CMPMODE::LE, (this->processDataNum+255)/256*256);  //256字节对齐生成mask,小于0为1
            //printf("%d, %d, %f, %f,%d\n",xLocal1(1),xLocal2(1), m_float1(1),m_float2(1),m_mask(1)); 
        }
        else if constexpr (std::is_same_v<DTYPE_X1, half>)
        {
            //printf( "half \n ");

            auto m_float1    = QueueFloat1.Get<float>();
            auto m_float2    = QueueFloat2.Get<float>();
            Cast(m_float1, xLocal1, RoundMode::CAST_NONE, this->processDataNum);
            Cast(m_float2, xLocal2, RoundMode::CAST_NONE, this->processDataNum);
            //1、input = |input - other| 
            Sub(m_float1, m_float1, m_float2, this->processDataNum);
            Abs(m_float1, m_float1, this->processDataNum);
            //2、other =  rtol * |other|
            Abs(m_float2, m_float2, this->processDataNum);
            Muls(m_float2, m_float2, this->rtol, this->processDataNum);
            //3、other = atol + rtol * |other|
            Adds(m_float2, m_float2, this->atol, this->processDataNum);
            //4、compare
            //DTYPE_Y scalar = 1.00;
            //Duplicate(m_zero, (DTYPE_X)(20.0), this->processDataNum);  
            Compare(m_mask, m_float1, m_float2, CMPMODE::LE, (this->processDataNum+255)/256*256);  //256字节对齐生成mask,小于0为1
            //printf("%f, %f, %f\n",m_float1(1),m_float2(1),m_mask(1));    
        }
        else if constexpr (std::is_same_v<DTYPE_X1, float>)
        {
            //printf( "DTYPE_X1 should be float, current type: %s\n ", typeid(DTYPE_X1).name());
            //printf("float \n");

            auto m_float1    = QueueFloat1.Get<float>();
            //1、input = |input - other| 
            Sub(m_float1, xLocal1, xLocal2, this->processDataNum);
            Abs(m_float1, m_float1, this->processDataNum);
            //2、other =  rtol * |other|
            Abs(xLocal2, xLocal2, this->processDataNum);
            Muls(xLocal2, xLocal2, this->rtol, this->processDataNum);
            //3、other = atol + rtol * |other|
            Adds(xLocal2, xLocal2, this->atol, this->processDataNum);
            //4、compare
            //DTYPE_Y scalar = 1.00;
            //Duplicate(m_zero, (DTYPE_X)(20.0), this->processDataNum); 
            Compare(m_mask, m_float1, xLocal2, CMPMODE::LE, (this->processDataNum+255)/256*256);  //256字节对齐生成mask,小于0为1
            //printf("%d, %d, %d\n",xLocal1(1),xLocal2(1),m_mask(1)); 
            
        }
        size_t m_size = (this->processDataNum+255)/256*256;
        //printf(" m_size is %d\n", m_size); 

        if constexpr (std::is_same_v<DTYPE_X1, float> || std::is_same_v<DTYPE_X1, half>) 
        {
            for (size_t i = 0; i < m_size; ++i) {
                uint8_t value = (m_mask.GetValue(i));
                //printf("v %d, %d\n",i, value); 
                for (size_t j = 0; j < 8; ++j) {
                    //printf("b %d %d ",i*8 + j, yLocal(i*8 + j)); 
                    yLocal(i*8 + j) =((value >> j) & 1);  
                    if((this->equal_nan == true) && isnan(xLocal1(i*8 + j)) && isnan(xLocal2(i*8 + j)))
                    {
                        yLocal(i*8 + j) = 1;
                        //printf("isnan(xLocal1(i*8 + j) %d\n",isnan(xLocal1(i*8 + j)));
                        //printf("isnan(xLocal2(i*8 + j) %d\n",isnan(xLocal2(i*8 + j)));
                        //printf(" a %d %d\n",i*8 + j, yLocal(i*8 + j)); 
                    }
                    //printf(" a %d %d\n",i*8 + j, yLocal(i*8 + j)); 
                    if(isnan(xLocal1(i*8 + j)) || isnan(xLocal1(i*8 + j)))
                    {
                        yLocal(i*8 + j) = 0;
                    }
                }
            }
        }
        else
        {
            for (size_t i = 0; i < m_size; ++i) {
                uint8_t value = (m_mask.GetValue(i));
                //printf("v %d, %d\n",i, value); 
                for (size_t j = 0; j < 8; ++j) {
                    //printf("b %d %d ",i*8 + j, yLocal(i*8 + j)); 
                    yLocal(i*8 + j) =((value >> j) & 1);  
                    //printf(" a %d %d\n",i*8 + j, yLocal(i*8 + j)); 
                }
            }
        }

        outQueueY.EnQue<DTYPE_Y>(yLocal);
        inQueueX1.FreeTensor(xLocal1);
        inQueueX2.FreeTensor(xLocal2);

    }
    __aicore__ inline void CopyOut(int32_t progress)
    {
        LocalTensor<DTYPE_Y> yLocal2 = outQueueY.DeQue<DTYPE_Y>();
        DataCopy(yGm[progress * this->tileDataNum], yLocal2, (this->processDataNum+32)/32*32);
        //("this->CoreDataNum %d, tileNum %d, tileDataNum %d TailDataNum %d\n",this->processDataNum, this->tileNum, this->tileDataNum, this->tailDataNum);
        //printf("this->processDataNum %d, progress %d, yGm1 %d yLocal1 %d\n",this->processDataNum, progress, yGm(2), yLocal2(2)); 
        outQueueY.FreeTensor(yLocal2);
    }   

private:
    TPipe pipe;
    //create queue for input, in this case depth is equal to buffer num
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX1;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX2;
    //create queue for output, in this case depth is equal to buffer num
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    GlobalTensor<DTYPE_X1> xGm1;
    GlobalTensor<DTYPE_X2> xGm2;
    GlobalTensor<DTYPE_Y> yGm;
    TBuf<QuePosition::VECCALC> QueueMask, QueueHalf1, QueueHalf2, QueueFloat1, QueueFloat2;

    //考生补充自定义成员变量
    uint32_t coreDataNum;
    uint32_t tileNum;
    uint32_t tileDataNum;
    uint32_t tailDataNum;
    uint32_t processDataNum;
    float rtol;
    float atol;
    bool equal_nan;
};

extern "C" __global__ __aicore__ void is_close(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelIsClose<DTYPE_X1,DTYPE_X2,DTYPE_Y> op;  
    //补充init和process函数调用内容
    op.Init(x1, x2, y, tiling_data.CoreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum, tiling_data.rtol, tiling_data.atol, tiling_data.equal_nan);
    op.Process();
}


